
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torch.optim as optim
import torchvision.datasets as datasets
import imageio
import numpy as np
import matplotlib
from torchvision.utils import make_grid, save_image
from torch.utils.data import DataLoader
from matplotlib import pyplot as plt
from tqdm import tqdm


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
cuda = torch.device('cuda:0')


# learning parameters
batch_size = 512
epochs = 100
sample_size = 64 # fixed sample size for generator
nz = 128 # latent vector size
k = 1 # number of steps to apply to the discriminator


transform = transforms.Compose([
                                transforms.ToTensor(),
                                transforms.Normalize((0.5,),(0.5,)),
])
to_pil_image = transforms.ToPILImage()

# Make input, output folders
!mkdir -p input
!mkdir -p outputs

# Load train data
train_data = datasets.MNIST(
    root='input/data',
    train=True,
    download=True,
    transform=transform
)
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)

A subdirectory or file -p already exists.
Error occurred while processing: -p.
A subdirectory or file input already exists.
Error occurred while processing: input.
A subdirectory or file -p already exists.
Error occurred while processing: -p.
A subdirectory or file outputs already exists.
Error occurred while processing: outputs.


class Generator(nn.Module):
    def __init__(self, nz):
        super(Generator, self).__init__()
        self.nz = nz
        self.main = nn.Sequential(
            nn.Linear(self.nz, 256),
            nn.LeakyReLU(0.2),
            nn.Linear(256, 512),
            nn.LeakyReLU(0.2),
            nn.Linear(512, 1024),
            nn.LeakyReLU(0.2),
            nn.Linear(1024, 784),
            nn.Tanh(),
        )
    def forward(self, x):
        return self.main(x).view(-1, 1, 28, 28)


class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.n_input = 784
        self.main = nn.Sequential(
            nn.Linear(self.n_input, 1024),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),
            nn.Linear(1024, 512),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),
            nn.Linear(256, 1),
            nn.Sigmoid(),
        )
    def forward(self, x):
        x = x.view(-1, 784)
        return self.main(x)


generator = Generator(nz).to(device)
discriminator = Discriminator().to(device)
print('##### GENERATOR #####')
print(generator)
print('######################')
print('\n##### DISCRIMINATOR #####')
print(discriminator)
print('######################')

##### GENERATOR #####
Generator(
  (main): Sequential(
    (0): Linear(in_features=128, out_features=256, bias=True)
    (1): LeakyReLU(negative_slope=0.2)
    (2): Linear(in_features=256, out_features=512, bias=True)
    (3): LeakyReLU(negative_slope=0.2)
    (4): Linear(in_features=512, out_features=1024, bias=True)
    (5): LeakyReLU(negative_slope=0.2)
    (6): Linear(in_features=1024, out_features=784, bias=True)
    (7): Tanh()
  )
)
######################

##### DISCRIMINATOR #####
Discriminator(
  (main): Sequential(
    (0): Linear(in_features=784, out_features=1024, bias=True)
    (1): LeakyReLU(negative_slope=0.2)
    (2): Dropout(p=0.3, inplace=False)
    (3): Linear(in_features=1024, out_features=512, bias=True)
    (4): LeakyReLU(negative_slope=0.2)
    (5): Dropout(p=0.3, inplace=False)
    (6): Linear(in_features=512, out_features=256, bias=True)
    (7): LeakyReLU(negative_slope=0.2)
    (8): Dropout(p=0.3, inplace=False)
    (9): Linear(in_features=256, out_features=1, bias=True)
    (10): Sigmoid()
  )
)
######################


# optimizers
optim_g = optim.Adam(generator.parameters(), lr=0.0002)
optim_d = optim.Adam(discriminator.parameters(), lr=0.0002)


# loss function
criterion = nn.BCELoss() # Binary Cross Entropy loss


losses_g = [] # to store generator loss after each epoch
losses_d = [] # to store discriminator loss after each epoch
images = [] # to store images generatd by the generator


# to create real labels (1s)
def label_real(size):
    data = torch.ones(size, 1)
    return data.to(device)
# to create fake labels (0s)
def label_fake(size):
    data = torch.zeros(size, 1)
    return data.to(device)


# function to create the noise vector
def create_noise(sample_size, nz):
    return torch.randn(sample_size, nz).to(device)


# to save the images generated by the generator
def save_generator_image(image, path):
    save_image(image, path)


# create the noise vector - fixed to track how GAN is trained.
noise = create_noise(sample_size, nz)


import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"


torch.manual_seed(7777)

def generator_loss(output, true_label):
    return criterion(output, true_label)

def discriminator_loss(output, true_label):
    return criterion(output, true_label)

# Normalize the images to [0, 1] range before visualization
def normalize(image):
    return (image - image.min()) / (image.max() - image.min())

for epoch in range(epochs):
    loss_g = 0.0
    loss_d = 0.0
    for bi, data in tqdm(enumerate(train_loader), total=int(len(train_data)/train_loader.batch_size)):
        image, _ = data
        image = image.to(device)
        b_size = image.size(0)

        # Train Discriminator
        optim_d.zero_grad()
        # Real Images
        real_label = label_real(b_size)
        output_real = discriminator(image)
        d_loss_real = discriminator_loss(output_real, real_label)

        # Fake Images
        noise = create_noise(b_size, nz)
        fake_images = generator(noise)
        fake_label = label_fake(b_size)
        output_fake = discriminator(fake_images.detach())
        d_loss_fake = discriminator_loss(output_fake, fake_label)

        # Total Discriminator Loss
        d_loss = d_loss_real + d_loss_fake
        d_loss.backward()
        optim_d.step()

        # Train Generator
        optim_g.zero_grad()
        output_fake = discriminator(fake_images)
        g_loss = generator_loss(output_fake, real_label)
        g_loss.backward()
        optim_g.step()

        loss_g += g_loss.item()
        loss_d += d_loss.item()

    # Create and visualize generated images every 5 epochs
    if (epoch + 1) % 5 == 0:
        generated_img = generator(noise).cpu().detach()
        generated_img = make_grid(generated_img)
        normalized_img = normalize(generated_img)  # Normalize the image
        plt.imshow(normalized_img.permute(1, 2, 0))
        plt.title(f'Epoch {epoch+1}')
        plt.axis('off')
        plt.show()

        # Save the generated torch tensor models to disk
        save_generator_image(generated_img, f"outputs/gen_img{epoch+1}.png")
        images.append(generated_img)

    epoch_loss_g = loss_g / bi # Total generator loss for the epoch
    epoch_loss_d = loss_d / bi # Total discriminator loss for the epoch
    losses_g.append(epoch_loss_g)
    losses_d.append(epoch_loss_d)
    
    print(f"Epoch {epoch+1} of {epochs}")
    print(f"Generator loss: {epoch_loss_g:.8f}, Discriminator loss: {epoch_loss_d:.8f}")

118it [00:10, 10.80it/s]

Epoch 1 of 100
Generator loss: 2.11920312, Discriminator loss: 0.76175348

118it [00:11, 10.26it/s]

Epoch 2 of 100
Generator loss: 3.21311511, Discriminator loss: 0.68023711

118it [00:11, 10.23it/s]

Epoch 3 of 100
Generator loss: 2.95189136, Discriminator loss: 0.69517191

118it [00:10, 11.05it/s]

Epoch 4 of 100
Generator loss: 2.54177409, Discriminator loss: 0.69475299

118it [00:11, 10.49it/s]

Epoch 5 of 100
Generator loss: 2.27097447, Discriminator loss: 0.71742772


print('DONE TRAINING')
torch.save(generator.state_dict(), 'outputs/generator.pth')

DONE TRAINING


# save the generated images as GIF file
imgs = [np.array(to_pil_image(img)) for img in images]
imageio.mimsave('outputs/generator_images.gif', imgs)


# plot and save the generator and discriminator loss
plt.figure()
plt.plot(losses_g, label='Generator loss')
plt.plot(losses_d, label='Discriminator Loss')
plt.legend()
plt.savefig('outputs/loss.png')


import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torch.optim as optim
import torchvision.datasets as datasets
from torchvision.utils import make_grid, save_image
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
from tqdm import tqdm

# Define learning parameters
batch_size = 512
epochs = 100
sample_size = 64  # fixed sample size for generator
nz = 128  # latent vector size
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Prepare training dataset
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,)),
])

train_data = datasets.MNIST(
    root='./input/data',
    train=True,
    download=True,
    transform=transform
)
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)

# Generator
class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.n_input = 784
        self.main = nn.Sequential(
            nn.Linear(self.n_input, 1024),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),  # adding dropout
            nn.Linear(1024, 512),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),  # adding dropout
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),  # adding dropout
            nn.Linear(256, 1),
            nn.Sigmoid(),
        )
    def forward(self, x):
        x = x.view(-1, 784)
        return self.main(x)

# Discriminator
class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.n_input = 784
        self.main = nn.Sequential(
            nn.Linear(self.n_input, 1024),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),
            nn.Linear(1024, 512),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),
            nn.Linear(256, 1),
            nn.Sigmoid(),
        )
    def forward(self, x):
        x = x.view(-1, 784)
        return self.main(x)

generator = Generator(nz).to(device)
discriminator = Discriminator().to(device)

# Optimizers and loss function
optim_g = optim.SGD(generator.parameters(), lr=0.0002)
optim_d = optim.Adam(discriminator.parameters(), lr=0.0002)
criterion = nn.BCELoss()

# Helper functions
def label_real(size):
    return torch.ones(size, 1, device=device)

def label_fake(size):
    return torch.zeros(size, 1, device=device)

def create_noise(sample_size, nz):
    return torch.randn(sample_size, nz).to(device)

def save_generator_image(image, path):
    save_image(image, path)

def normalize(image):
    return (image - image.min()) / (image.max() - image.min())

noise = create_noise(sample_size, nz)

# Training Loop
torch.manual_seed(7777)

def generator_loss(disc_fake_pred):
    gen_loss = torch.mean(torch.log(1 - disc_fake_pred))
    return gen_loss

def discriminator_loss(real_pred, fake_pred):
    real_loss = criterion(real_pred, torch.ones_like(real_pred))
    fake_loss = criterion(fake_pred, torch.zeros_like(fake_pred))
    total_loss = (real_loss + fake_loss) / 2
    return total_loss

# Initialize lists to store the losses for plotting
g_losses = []
d_losses = []

# Store images from specified epochs
epoch_images = {}

for epoch in range(epochs):
    loss_g = 0.0
    loss_d = 0.0
    for bi, data in tqdm(enumerate(train_loader), total=int(len(train_data)/train_loader.batch_size)):
        image, _ = data
        image = image.to(device)
        b_size = image.size(0)

        # Train Discriminator
        optim_d.zero_grad()
        real_label = label_real(b_size)
        output_real = discriminator(image)
        noise = create_noise(b_size, nz)
        fake_images = generator(noise)
        output_fake = discriminator(fake_images.detach())
        d_loss = discriminator_loss(output_real, output_fake)
        d_loss.backward()
        optim_d.step()

        # Train Generator
        optim_g.zero_grad()
        output_fake_for_gen = discriminator(fake_images)
        # Change generator loss to maximize log(D(G(z)))
        g_loss = criterion(output_fake_for_gen, label_real(b_size))
        g_loss.backward()
        optim_g.step()

        loss_g += g_loss.item()
        loss_d += d_loss.item()

    g_losses.append(loss_g / len(train_loader))
    d_losses.append(loss_d / len(train_loader))

    # Save images at specified epochs
    if epoch in [0, 49, 99]:  # Epochs are zero-indexed
        with torch.no_grad():
            generated_img = generator(noise).detach().cpu()
            generated_img = make_grid(generated_img, normalize=True)
            save_image(generated_img, f"./outputs/gen_img_epoch{epoch+1}.png")
            epoch_images

118it [00:10, 10.83it/s]                         
118it [00:09, 11.84it/s]                         
118it [00:10, 10.77it/s]                         
118it [00:10, 10.92it/s]                         
118it [00:11, 10.69it/s]                         
118it [00:11, 10.71it/s]                         
118it [00:11, 10.49it/s]                         
118it [00:11, 10.68it/s]                         
118it [00:10, 10.74it/s]                         
118it [00:11, 10.61it/s]                         
118it [00:11, 10.69it/s]                         
118it [00:11, 10.61it/s]                         
118it [00:11, 10.55it/s]                         
118it [00:10, 11.34it/s]                         
118it [00:11, 10.54it/s]                         
118it [00:10, 10.80it/s]                         
118it [00:11, 10.51it/s]                         
118it [00:10, 10.83it/s]                         
118it [00:11, 10.63it/s]                         
118it [00:11, 10.65it/s]                         
118it [00:10, 10.74it/s]                         
118it [00:11, 10.56it/s]                         
118it [00:11, 10.62it/s]                         
118it [00:10, 10.78it/s]                         
118it [00:11, 10.63it/s]                         
118it [00:11, 10.66it/s]                         
118it [00:11, 10.66it/s]                         
118it [00:10, 11.01it/s]                         
118it [00:10, 10.88it/s]                         
118it [00:10, 11.53it/s]                         
118it [00:11, 10.56it/s]                         
118it [00:11, 10.67it/s]                         
118it [00:11, 10.56it/s]                         
118it [00:10, 10.78it/s]                         
118it [00:11, 10.68it/s]                         
118it [00:10, 10.75it/s]                         
118it [00:11, 10.68it/s]                         
118it [00:11, 10.61it/s]                         
118it [00:11, 10.67it/s]                         
118it [00:11, 10.63it/s]                         
118it [00:11, 10.71it/s]                         
118it [00:10, 10.85it/s]                         
118it [00:10, 10.76it/s]                         
118it [00:11, 10.64it/s]                         
118it [00:11, 10.66it/s]                         
118it [00:11, 10.53it/s]                         
118it [00:10, 10.75it/s]                         
118it [00:10, 10.73it/s]                         
118it [00:11, 10.57it/s]                         
118it [00:11, 10.57it/s]                         
118it [00:11, 10.71it/s]                         
118it [00:10, 10.75it/s]                         
118it [00:11, 10.67it/s]                         
118it [00:11, 10.69it/s]                         
118it [00:11, 10.56it/s]                         
118it [00:11, 10.53it/s]                         
118it [00:11, 10.63it/s]                         
118it [00:10, 10.75it/s]                         
118it [00:10, 10.80it/s]                         
118it [00:11, 10.63it/s]                         
118it [00:11, 10.71it/s]                         
118it [00:11, 10.67it/s]                         
118it [00:11, 10.57it/s]                         
118it [00:10, 10.84it/s]                         
118it [00:11, 10.64it/s]                         
118it [00:11, 10.64it/s]                         
118it [00:11, 10.69it/s]                         
118it [00:11, 10.70it/s]                         
118it [00:11, 10.59it/s]                         
118it [00:11, 10.69it/s]                         
118it [00:11, 10.71it/s]                         
118it [00:11, 10.61it/s]                         
118it [00:10, 10.77it/s]                         
118it [00:10, 11.29it/s]                         
118it [00:11, 10.37it/s]                         
118it [00:10, 10.77it/s]                         
118it [00:11, 10.66it/s]                         
118it [00:11, 10.71it/s]                         
118it [00:10, 11.36it/s]                         
118it [00:11, 10.66it/s]                         
118it [00:11, 10.70it/s]                         
118it [00:11, 10.70it/s]                         
118it [00:10, 10.89it/s]                         
118it [00:11, 10.59it/s]                         
118it [00:11, 10.73it/s]                         
118it [00:11, 10.69it/s]                         
118it [00:11, 10.69it/s]                         
118it [00:11, 10.60it/s]                         
118it [00:10, 11.26it/s]                         
118it [00:11, 10.67it/s]                         
118it [00:10, 10.85it/s]                         
118it [00:11, 10.52it/s]                         
118it [00:10, 11.02it/s]                         
118it [00:10, 11.02it/s]                         
118it [00:11, 10.42it/s]                         
118it [00:11, 10.65it/s]                         
118it [00:10, 11.02it/s]                         
118it [00:10, 10.74it/s]                         
118it [00:11, 10.63it/s]                         
118it [00:10, 10.99it/s]


print('DONE TRAINING')
torch.save(generator.state_dict(), 'outputs/generator_mod.pth')

DONE TRAINING


# Ring a sound when this cell finishes running
import os
os.system('say "Training complete"')

1


# Show the saved images for epoch 1, 50, and 100
for epoch in [1, 50, 100]:  # Corrected to match epoch numbers starting from 1
    image_file = f"./outputs/gen_img_epoch{epoch}.png"
    image = plt.imread(image_file)
    plt.figure(figsize=(8, 8))
    plt.axis("off")
    plt.title(f"Generated Images at Epoch {epoch}")
    plt.imshow(image)
    plt.show()


# Plot the training losses
plt.figure(figsize=(10, 5))
plt.title("Generator and Discriminator Loss During Training")
plt.plot(g_losses, label="Generator")
plt.plot(d_losses, label="Discriminator")
plt.xlabel("Epochs")
plt.ylabel("Loss")
plt.legend()
plt.show()


import torch
from torch import nn, autograd
from torch.utils.data import DataLoader
from torchvision import transforms, datasets
from torchvision.utils import save_image, make_grid
import numpy as np
import torch.optim as optim
import os
from tqdm import tqdm
import matplotlib.pyplot as plt

# Spectral Normalization code from PyTorch
def spectral_norm(module, name='weight', n_power_iterations=1, eps=1e-12, dim=None):
    if dim is None:
        if isinstance(module, (nn.ConvTranspose2d, nn.Conv2d)):
            dim = 1
        else:
            dim = 0
    setattr(module, name, nn.utils.spectral_norm(
        getattr(module, name), n_power_iterations=n_power_iterations, eps=eps, dim=dim))

# Improved Discriminator with Spectral Normalization
class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.n_input = 784
        self.main = nn.Sequential(
            nn.utils.spectral_norm(nn.Linear(self.n_input, 1024)),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),
            nn.utils.spectral_norm(nn.Linear(1024, 512)),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),
            nn.utils.spectral_norm(nn.Linear(512, 256)),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),
            nn.utils.spectral_norm(nn.Linear(256, 1)),
        )

    def forward(self, x):
        x = x.view(-1, 784)
        return self.main(x)

    
class Generator(nn.Module):
    def __init__(self, nz):
        super(Generator, self).__init__()
        self.nz = nz
        self.main = nn.Sequential(
            nn.Linear(self.nz, 256),
            nn.LeakyReLU(0.2),
            nn.Linear(256, 512),
            nn.LeakyReLU(0.2),
            nn.Linear(512, 1024),
            nn.LeakyReLU(0.2),
            nn.Linear(1024, 784),
            nn.Tanh(),
        )
    def forward(self, x):
        return self.main(x).view(-1, 1, 28, 28)
    
# WGAN-GP Loss and Gradient Penalty
def gradient_penalty(discriminator, real_data, fake_data, device):
    batch_size = real_data.size(0)
    alpha = torch.rand(batch_size, 1, 1, 1).to(device)
    alpha = alpha.expand_as(real_data)
    
    interpolated = alpha * real_data.data + (1 - alpha) * fake_data.data
    interpolated = interpolated.requires_grad_(True)
    
    prob_interpolated = discriminator(interpolated)
    
    gradients = autograd.grad(outputs=prob_interpolated, inputs=interpolated,
                              grad_outputs=torch.ones(prob_interpolated.size()).to(device),
                              create_graph=True, retain_graph=True)[0]
    
    gradients = gradients.view(batch_size, -1)
    gradients_norm = torch.sqrt(torch.sum(gradients ** 2, dim=1) + 1e-12)
    
    return torch.mean((gradients_norm - 1) ** 2)


# Hyperparameters
lambda_gp = 10  # Gradient penalty lambda hyperparameter
n_critic = 5  # The number of critic iterations for one step of the generator
clip_value = 0.01  # Gradient clipping value
nz = 128  # nz is the length of the noise vector
batch_size = 512
epochs = 100
sample_size = 64 # fixed sample size for generator
k = 1 # number of steps to apply to the discriminator
lr_d = 0.0001  # learning rate for the discriminator
lr_g = 0.0002  # learning rate for the generator
discriminator_noise = 0.05  # Noise to add to discriminator inputs

# Now when you instantiate the Generator and Discriminator, make sure to pass `nz` to the Generator
generator = Generator(nz).to(device)
discriminator = Discriminator().to(device)


transform = transforms.Compose([
                                transforms.ToTensor(),
                                transforms.Normalize((0.5,),(0.5,)),
])
to_pil_image = transforms.ToPILImage()

# Make input, output folders
!mkdir -p input
!mkdir -p outputs

# Load train data
train_data = datasets.MNIST(
    root='input/data',
    train=True,
    download=True,
    transform=transform
)
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)

A subdirectory or file -p already exists.
Error occurred while processing: -p.
A subdirectory or file input already exists.
Error occurred while processing: input.
A subdirectory or file -p already exists.
Error occurred while processing: -p.
A subdirectory or file outputs already exists.
Error occurred while processing: outputs.


# optimizers
optim_g = optim.Adam(generator.parameters(), lr=lr_g, betas=(0.5, 0.999))
optim_d = optim.Adam(discriminator.parameters(), lr=lr_d, betas=(0.5, 0.999))

# Initialize learning rate scheduler
scheduler_g = torch.optim.lr_scheduler.StepLR(optim_g, step_size=1000, gamma=0.97)
scheduler_d = torch.optim.lr_scheduler.StepLR(optim_d, step_size=1000, gamma=0.97)


# loss function
criterion = nn.BCELoss() # Binary Cross Entropy loss


losses_g = [] # to store generator loss after each epoch
losses_d = [] # to store discriminator loss after each epoch
images = [] # to store images generatd by the generator


# to create real labels (1s)
def label_real(size):
    data = torch.ones(size, 1)
    return data.to(device)
# to create fake labels (0s)
def label_fake(size):
    data = torch.zeros(size, 1)
    return data.to(device)

# function to create the noise vector
def create_noise(sample_size, nz):
    return torch.randn(sample_size, nz).to(device)

# to save the images generated by the generator
def save_generator_image(image, path):
    save_image(image, path)
    
# create the noise vector - fixed to track how GAN is trained.
noise = create_noise(sample_size, nz)


import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"


torch.manual_seed(7777)

# Wasserstein loss function for the generator
def generator_loss(output):
    return -output.mean()

# Wasserstein loss function for the discriminator
def discriminator_loss(output_real, output_fake):
    return output_fake.mean() - output_real.mean()

# Fixed noise for tracking generator training
fixed_noise = create_noise(sample_size, nz)

# Training Loop
for epoch in range(epochs):
    loss_g = 0.0
    loss_d = 0.0
    for bi, data in tqdm(enumerate(train_loader), total=int(len(train_data)/train_loader.batch_size)):
        image, _ = data
        image = image.to(device)
        b_size = image.size(0)
        noise_data = create_noise(b_size, nz)

        ### Update Discriminator ###
        for _ in range(n_critic):
            optim_d.zero_grad()

            # Discriminator on real data
            real_images = image
            output_real = discriminator(real_images)
            d_loss_real = -output_real.mean()

            # Discriminator on fake data
            fake_images = generator(noise_data).detach()
            output_fake = discriminator(fake_images)
            d_loss_fake = output_fake.mean()

            # Gradient penalty
            # Make sure gradient_penalty function is defined elsewhere in your code
            penalty = gradient_penalty(discriminator, real_images, fake_images, device)
            d_loss = d_loss_real + d_loss_fake + lambda_gp * penalty

            d_loss.backward()
            optim_d.step()

        ### Update Generator ###
        optim_g.zero_grad()
        fake_images = generator(noise_data)
        output_fake = discriminator(fake_images)
        g_loss = -output_fake.mean()
        g_loss.backward()
        optim_g.step()

        loss_g += g_loss.item()
        loss_d += d_loss.item()

    # Update learning rate if you're using a learning rate scheduler
    # Make sure scheduler_g and scheduler_d are defined elsewhere in your code
    scheduler_g.step()
    scheduler_d.step()

    # Generate and save images every 5 epochs
    # ... [image generation and saving code]

    # Log the losses for this epoch
    epoch_loss_g = loss_g / len(train_loader)
    epoch_loss_d = loss_d / (len(train_loader) * n_critic)
    losses_g.append(epoch_loss_g)
    losses_d.append(epoch_loss_d)

    print(f"Epoch {epoch+1} of {epochs}")
    print(f"Generator loss: {epoch_loss_g:.8f}, Discriminator loss: {epoch_loss_d:.8f}")
    
    # Create and visualize generated images every 5 epochs
    if (epoch + 1) % 5 == 0:
        generated_img = generator(noise).cpu().detach()
        generated_img = make_grid(generated_img)
        plt.imshow(generated_img.permute(1, 2, 0))
        plt.title(f'Epoch {epoch+1}')
        plt.axis('off')
        plt.show()

        # Save the generated torch tensor models to disk
        save_generator_image(generated_img, f"outputs/gen_img{epoch+1}.png")
        images.append(generated_img)

118it [00:16,  7.20it/s]

Epoch 1 of 100
Generator loss: -8.26281330, Discriminator loss: -0.46020975

118it [00:16,  7.35it/s]

Epoch 2 of 100
Generator loss: -8.60957666, Discriminator loss: -0.46775316

118it [00:16,  7.09it/s]

Epoch 3 of 100
Generator loss: -8.54908839, Discriminator loss: -0.46157440

118it [00:16,  7.12it/s]

Epoch 4 of 100
Generator loss: -8.42119303, Discriminator loss: -0.48207996

118it [00:17,  6.82it/s]                         
Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).

Epoch 5 of 100
Generator loss: -7.82988271, Discriminator loss: -0.47474676


print('DONE TRAINING')
torch.save(generator.state_dict(), 'outputs/generator_improved.pth')

DONE TRAINING


import imageio


# save the generated images as GIF file
imgs = [np.array(to_pil_image(img)) for img in images]
imageio.mimsave('outputs/generator_images_improved.gif', imgs)


# After training loop ends
# Plot the training losses
plt.figure(figsize=(10,5))
plt.title("Generator and Discriminator Loss During Training")
plt.plot(losses_g,label="G")
plt.plot(losses_d,label="D")
plt.xlabel("iterations")
plt.ylabel("Loss")
plt.legend()
plt.show()

# Save models at the end of training
torch.save(generator.state_dict(), 'generator_improved.pth')
torch.save(discriminator.state_dict(), 'discriminator_improved.pth')

Homework 6 Code - Dario Placencio¶

1 Implementation: GAN (50 pts)¶

GAN with MNIST dataset¶

Define learning parameters¶

Prepare training dataset¶

Generator¶

Discriminator¶

Tools for training¶

Training¶

Variation of GAN - SGD¶

GAN Improvements¶